# Testing molct model.
# 23-11-01

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# setting context & hyper parameters
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

from cybertron.model.molct import MolCT
from ms_cybertron.model import MolCTWithEdge as ms_MolCT

# generate data for test
np.random.seed(42)
node_vec = np.random.uniform(size=(3, 5, 4)) # (B, A, F)
node_mask = np.sum(node_vec, axis=-1) < 2 # (B, A)

edge_vec = np.random.uniform(size=(3, 5, 5, 4)) # (B, A, A, F)
eye = np.eye(5).reshape(1, 5, 5, 1)
eye_mask = np.array(eye, dtype=np.bool_)
edge_vec = np.where(eye_mask, eye, edge_vec)

edge_mask = np.sum(edge_vec, axis=-1) < 2 # (B, A, A)
eye = np.eye(5).reshape(1, 5, 5)
eye_mask = np.array(eye, dtype=np.bool_)
edge_mask = np.where(eye_mask, eye_mask, edge_mask)

edge_cutoff = np.random.uniform(size=(3, 5, 5)) # (B, A, A)
edge_cutoff = np.where(eye_mask, eye, edge_cutoff)

is_edge_update = True
is_coupled_interaction = False
num_interaction = 3
print("=================Test MolCT====================")

def molct_fn(a, b, c, d, e):
    out = MolCT(dim_feature=4,
                dim_node_emb=4,
                dim_edge_emb=4,
                is_edge_update=is_edge_update,
                is_coupled_interaction=is_coupled_interaction,
                n_heads=2,
                dim_outer_pdct=8,
                num_transition=2,
                n_interaction=num_interaction,
                ) # type: ignore
    return out(a, b, c, d, e)

molct_fn = jax.vmap(molct_fn, in_axes=(0, 0, 0, 0, 0))
ts_molct_fn = hk.transform(molct_fn, apply_rng=True)
molct_params = ts_molct_fn.init(rng, node_vec, node_mask, edge_vec, edge_mask, edge_cutoff)
node_out, edge_out = jit(ts_molct_fn.apply)(molct_params, rng, node_vec, node_mask, edge_vec, edge_mask, edge_cutoff)
print(f"node_out: shape of {node_out.shape}, value of \n{node_out}")
print(f"edge_out: shape of {edge_out.shape}, value of \n{edge_out}")

ms_molct_fn = ms_MolCT(dim_feature=4,
                       n_interaction=num_interaction,
                       edge_update=is_edge_update,
                       n_heads=2,
                       fixed_cycles=False,
                       coupled_interaction=is_coupled_interaction,
                       dim_outer_product=8,
                       num_intermediate_factor=2,
                       )
ms_molct_fn.set_dimension(4, 4)

ms_params = {}
if is_coupled_interaction:
    _name = 'molct/~build_interaction/interaction_unit/~/'
    _params = hk.data_structures.to_immutable_dict(molct_params)
    _ms_name = 'interaction.0'
    ms_params[f'{_ms_name}.positional_embedding.norm.gamma'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/norm']['scale']), ms.float32))
    ms_params[f'{_ms_name}.positional_embedding.norm.beta'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/norm']['offset']), ms.float32))
    ms_params[f'{_ms_name}.positional_embedding.x2q.weight'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/q_gen']['w']).T, ms.float32))
    ms_params[f'{_ms_name}.positional_embedding.x2k.weight'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/k_gen']['w']).T, ms.float32))
    ms_params[f'{_ms_name}.positional_embedding.x2v.weight'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/v_gen']['w']).T, ms.float32))
    ms_params[f'{_ms_name}.multi_head_attention.output.weight'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}multihead_attention/linear_output']['w']).T, ms.float32))
    ms_params[f'{_ms_name}.norm.beta'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}node_norm']['offset']), ms.float32))
    ms_params[f'{_ms_name}.norm.gamma'] = \
        ms.Parameter(Tensor(np.array(_params[f'{_name}node_norm']['scale']), ms.float32))
    if is_edge_update:
        ms_params[f'{_ms_name}.edge_norm.beta'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}edge_norm']['offset']), ms.float32))
        ms_params[f'{_ms_name}.edge_norm.gamma'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}edge_norm']['scale']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.left_projection_weights'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['left_weights']).T, ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.right_projection_weights'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['right_weights']).T, ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.left_projection_biases'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['left_bias']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.right_projection_biases'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['right_bias']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.linear_output_weights'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['output_weights']).T, ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.o_biases'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['output_bias']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.layer_norm.beta'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product/~/norm_fn']['offset']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.layer_norm.gamma'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product/~/norm_fn']['scale']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition1_weights'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_1']['w']).T, ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition1_biases'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_1']['b']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition2_weights'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_2']['w']).T, ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition2_biases'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_2']['b']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.layer_norm.beta'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/norm']['offset']), ms.float32))
        ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.layer_norm.gamma'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/norm']['scale']), ms.float32))

else:
    _params = hk.data_structures.to_immutable_dict(molct_params)
    for _ in range(num_interaction):
        _name = f'molct/~build_interaction/interaction_unit_{_}/~/'
    
        _ms_name = f'interaction.{_}'
        ms_params[f'{_ms_name}.positional_embedding.norm.gamma'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/norm']['scale']), ms.float32))
        ms_params[f'{_ms_name}.positional_embedding.norm.beta'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/norm']['offset']), ms.float32))
        ms_params[f'{_ms_name}.positional_embedding.x2q.weight'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/q_gen']['w']).T, ms.float32))
        ms_params[f'{_ms_name}.positional_embedding.x2k.weight'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/k_gen']['w']).T, ms.float32))
        ms_params[f'{_ms_name}.positional_embedding.x2v.weight'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}positional_embedding/v_gen']['w']).T, ms.float32))
        ms_params[f'{_ms_name}.multi_head_attention.output.weight'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}multihead_attention/linear_output']['w']).T, ms.float32))
        ms_params[f'{_ms_name}.norm.beta'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}node_norm']['offset']), ms.float32))
        ms_params[f'{_ms_name}.norm.gamma'] = \
            ms.Parameter(Tensor(np.array(_params[f'{_name}node_norm']['scale']), ms.float32))
        if is_edge_update:
            ms_params[f'{_ms_name}.edge_norm.beta'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}edge_norm']['offset']), ms.float32))
            ms_params[f'{_ms_name}.edge_norm.gamma'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}edge_norm']['scale']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.left_projection_weights'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['left_weights']).T, ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.right_projection_weights'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['right_weights']).T, ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.left_projection_biases'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['left_bias']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.right_projection_biases'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['right_bias']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.linear_output_weights'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['output_weights']).T, ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.o_biases'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product']['output_bias']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.layer_norm.beta'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product/~/norm_fn']['offset']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.outer_product.layer_norm.gamma'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/outer_product/~/norm_fn']['scale']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition1_weights'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_1']['w']).T, ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition1_biases'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_1']['b']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition2_weights'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_2']['w']).T, ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.transition2_biases'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/transition_2']['b']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.layer_norm.beta'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/norm']['offset']), ms.float32))
            ms_params[f'{_ms_name}.hyperformer_pair_block.pair_transition.layer_norm.gamma'] = \
                ms.Parameter(Tensor(np.array(_params[f'{_name}hyperformer_pair_block/transition/~/norm']['scale']), ms.float32))

load_param_into_net(ms_molct_fn, ms_params)

ms_node_vec = Tensor(node_vec, ms.float32)
ms_node_mask = Tensor(node_mask, ms.bool_)
ms_edge_vec = Tensor(edge_vec, ms.float32)
ms_edge_mask = Tensor(edge_mask, ms.bool_)
ms_edge_cutoff = Tensor(edge_cutoff, ms.float32)

ms_node_out, ms_edge_out = ms_molct_fn(ms_node_vec, ms_node_mask, ms_edge_vec, ms_edge_mask, ms_edge_cutoff) # type: ignore
print(f"ms_node_out: shape of {ms_node_out.shape}, value of \n{ms_node_out}")
print(f"ms_edge_out: shape of {ms_edge_out.shape}, value of \n{ms_edge_out}")

node_same = np.allclose(node_out, ms_node_out.asnumpy(), atol=EPSILON)
edge_same = np.allclose(edge_out, ms_edge_out.asnumpy(), atol=EPSILON)
print(f"node same: {node_same} in accuracy of {EPSILON}")
print(f"edge same: {edge_same} in accuracy of {EPSILON}")

node_diff = node_out - ms_node_out.asnumpy()
node_norm = np.linalg.norm(node_diff)
edge_diff = edge_out - ms_edge_out.asnumpy()
edge_norm = np.linalg.norm(edge_diff)
print(f"node norm: {node_norm}")
print(f"edge norm: {edge_norm}")
